﻿using System;
using System.Collections.Generic;
using System.Globalization;
using System.Linq;
using System.Diagnostics;
using EvoBoost.Core.Evolutionary;

namespace EvoBoost.Core
{
	public class Booster<T>
	{
		public bool PrintVerboseOutput { get; set; }

		#region Boosting parameters

		private ILossFunction proxyLossFunction;
		public ILossFunction ProxyLossFunction
		{
			get { return this.proxyLossFunction; }
			set
			{
				if (value == null)
					throw new ArgumentNullException("value");

				this.proxyLossFunction = value;
			}
		}

		private ILossFunction lossFunction;
		public ILossFunction LossFunction
		{
			get { return this.lossFunction; }
			set
			{
				if (value == null)
					throw new ArgumentNullException("value");

				this.lossFunction = value;
			}
		}

		private int boostingIterations;
		public int BoostingIterations
		{
			get { return this.boostingIterations; }
			set
			{
				if (value <= 0)
					throw new ArgumentOutOfRangeException("value", "Iteration count should be positive.");

				this.boostingIterations = value;
			}
		}

		private int minSamplesInNode;
		public int MinSamplesInNode
		{
			get { return this.minSamplesInNode; }
			set
			{
				if (value <= 0)
					throw new ArgumentOutOfRangeException("value", "Min number of samples in node should be positive.");

				this.minSamplesInNode = value;
			}
		}

		private int maxTreeRegions;
		public int MaxTreeRegions
		{
			get { return this.maxTreeRegions; }
			set
			{
				if (value <= 1)
					throw new ArgumentOutOfRangeException("value", "Max number of terminal nodes should be greater than 1.");

				this.maxTreeRegions = value;
			}
		}

		private double shrinkage;
		public double Shrinkage
		{
			get { return this.shrinkage; }
			set
			{
				if (value <= 0.0 || value > 1.0)
					throw new ArgumentOutOfRangeException("value", "Shrinkage coefficient must belong to (0,1] range.");

				this.shrinkage = value;
			}
		}

		private double proxyLossTrimmingRate;
		public double ProxyLossTrimmingRate
		{
			get { return this.proxyLossTrimmingRate; }
			set
			{
				if (this.proxyLossTrimmingRate < 0 || this.proxyLossTrimmingRate >= 0.9)
					throw new ArgumentOutOfRangeException("value", "Proxy loss trimming rate should be in [0, 0.9] range.");
				this.proxyLossTrimmingRate = value;
			}
		}

		private int messageLevel;
		public int MessageLevel
		{
			get { return this.messageLevel; }
			set
			{
				if (value < 0)
					throw new ArgumentOutOfRangeException("value", "Message level should be non-negative.");

				this.messageLevel = value;
			}
		}

		#endregion

		#region Genetic algorithm parameters

		private double adjustSplitterMutationProbability;
		public double AdjustSplitterMutationProbability
		{
			get { return this.adjustSplitterMutationProbability; }
			set
			{
				if (value < 0 || value > 1)
					throw new ArgumentOutOfRangeException("value", "'Adjust splitter' mutation probability should belong to [0,1] range.");
				this.adjustSplitterMutationProbability = value;
			}
		}

		private int populationSize;
		public int PopulationSize
		{
			get { return this.populationSize; }
			set
			{
				if (value <= 0)
					throw new ArgumentOutOfRangeException("value", "Population size should be positive.");
				this.populationSize = value;
			}
		}

		private int maxEpochCount;
		public int MaxEpochCount
		{
			get { return this.maxEpochCount; }
			set
			{
				if (value <= 0)
					throw new ArgumentOutOfRangeException("value", "Max epoch count should be positive.");
				this.maxEpochCount = value;
			}
		}

		private int maxEpochCountWithoutImprovement;
		public int MaxEpochCountWithoutImprovement
		{
			get { return this.maxEpochCountWithoutImprovement; }
			set
			{
				if (value <= 0)
					throw new ArgumentOutOfRangeException("value", "Max epoch without improve count should be positive.");
				this.maxEpochCountWithoutImprovement = value;
			}
		}

		private double mutationRate;
		public double MutationRate
		{
			get { return this.mutationRate; }
			set
			{
				if (value <= 0 || value >= 1)
					throw new ArgumentOutOfRangeException("value", "Mutation rate should lie between 0 and 1.");
				this.mutationRate = value;
			}
		}

		private double crossoverRate;
		public double CrossoverRate
		{
			get { return this.crossoverRate; }
			set
			{
				if (value <= 0 || value >= 1)
					throw new ArgumentOutOfRangeException("value", "Crossover rate should lie between 0 and 1.");
				this.crossoverRate = value;
			}
		}

		#endregion

		public Booster()
		{
			this.proxyLossFunction = new SquaredLossFunction();
			this.lossFunction = new SquaredLossFunction();
			this.messageLevel = 1;
			this.shrinkage = 1;
			this.boostingIterations = 100;
			this.maxTreeRegions = 8;
			this.minSamplesInNode = 10;
			this.populationSize = 100;
			this.maxEpochCount = 30;
			this.maxEpochCountWithoutImprovement = 10;
			this.mutationRate = 0.2;
			this.crossoverRate = 0.7;
			this.adjustSplitterMutationProbability = 0.5;
			this.proxyLossTrimmingRate = 0;
			this.PrintVerboseOutput = true;
		}

		private void PrintIfVerbose(int level, string format, params object[] arg)
		{
			if (this.PrintVerboseOutput && level <= this.messageLevel)
				Console.Write(format, arg);
		}

		private void PrintLineIfVerbose(int level, string format, params object[] arg)
		{
			if (this.PrintVerboseOutput && level <= this.messageLevel)
				Console.WriteLine(format, arg);
		}

		public Forest<T> Boost(ICollection<Sample<T>> trainingSet, SplitBase<T> split)
		{
			if (trainingSet == null) throw new ArgumentNullException("trainingSet");
			if (split == null) throw new ArgumentNullException("split");
			if (trainingSet.Count < this.minSamplesInNode)
				throw new ArgumentException(
					String.Format(CultureInfo.CurrentCulture, "At least {0} training samples should be provided.", this.minSamplesInNode),
					"trainingSet");

			// Output boosting configuration to console
			PrintConfiguration();

			// Create forest with only the base tree
			Tree<T> baseTree = CreateBaseTree(trainingSet);
			Forest<T> boostedForest = new Forest<T>(baseTree, new List<Tree<T>>(), this.shrinkage);

			// Create holder for reduced training set
			List<Sample<T>> reducedTrainingSet = new List<Sample<T>>(trainingSet);

			// Create holder for proxy loss values
			List<double> proxyLossValues = new List<double>();
			List<double> sortedProxyLossValues = new List<double>();

			// Remember start time
			DateTime startTime = DateTime.Now;

			// Perform iterational boosting
			for (int iteration = 1; iteration <= this.boostingIterations; ++iteration)
			{
				if (reducedTrainingSet.Count < this.minSamplesInNode)
				{
					this.PrintLineIfVerbose(1, "Boosting convergence detected.");
					break;
				}

				// Start iteration
				DateTime iterationStartTime = DateTime.Now;
				this.PrintLineIfVerbose(1, "Iteration {0} started", iteration);

				// Train on reduced set
				boostedForest = PerformBoostingIteration(reducedTrainingSet, split, boostedForest);

				// Recalculate total loss and store loss values for loss trimming threshold selection
				if (this.proxyLossTrimmingRate > 0)
					proxyLossValues.Clear();
				double totalLoss = 0, totalProxyLoss = 0;
				foreach (Sample<T> sample in trainingSet)
				{
					double proxyLoss = this.proxyLossFunction.Calculate(
						sample.Outcome, boostedForest.CalculateOutcome(sample.Properties));
					double loss = this.lossFunction.Calculate(
						sample.Outcome, boostedForest.CalculateOutcome(sample.Properties));
					
					totalLoss += loss;
					totalProxyLoss += proxyLoss;

					proxyLossValues.Add(proxyLoss);
				}

				// Select threshold for loss trimming
				double lossThreshold = 0;
				if (this.proxyLossTrimmingRate > 0)
				{
					sortedProxyLossValues.Clear();
					sortedProxyLossValues.AddRange(proxyLossValues);
					sortedProxyLossValues.Sort();

					double lossSum = 0;
					int index = 0;
					while (index < sortedProxyLossValues.Count - 1 && lossSum <= this.proxyLossTrimmingRate * totalProxyLoss)
						lossSum += sortedProxyLossValues[index++];
					lossThreshold = sortedProxyLossValues[index];
				}
				
				// Update reduced training set
				reducedTrainingSet.Clear();
				int sampleIndex = 0;
				foreach (Sample<T> sample in trainingSet)
				{
					if (proxyLossValues[sampleIndex++] >= lossThreshold)
						reducedTrainingSet.Add(sample);
				}

				// Average loss
				double meanProxyLoss = totalProxyLoss / trainingSet.Count;
				double meanLoss = totalLoss / trainingSet.Count;

				// Finish iteration
				DateTime iterationEndTime = DateTime.Now;
				TimeSpan iterationTime = iterationEndTime - iterationStartTime;
				TimeSpan totalTime = iterationEndTime - startTime;

				// Print statistics
				this.PrintLineIfVerbose(1, "Iteration {0} finished in {1}", iteration, iterationTime);
				this.PrintLineIfVerbose(1, "Total training time is {0}", totalTime);
				this.PrintLineIfVerbose(
					1, "Average time per iteration is {0}", TimeSpan.FromSeconds(totalTime.TotalSeconds / iteration));
				this.PrintLineIfVerbose(1, "Average proxy loss: {0:0.0000}", meanProxyLoss);
				this.PrintLineIfVerbose(1, "Average loss: {0:0.0000}", meanLoss);
				if (this.proxyLossTrimmingRate > 0)
				{
					this.PrintLineIfVerbose(
						1, "{0} of {1} samples will be used in next iteration", reducedTrainingSet.Count, trainingSet.Count);
				}
				this.PrintLineIfVerbose(2, boostedForest.Trees[boostedForest.Trees.Count - 1].ToString());
				this.PrintLineIfVerbose(1, "");
			}

			return boostedForest;
		}

		private void PrintConfiguration()
		{
			this.PrintLineIfVerbose(1, "Booster configuration: ");
			this.PrintLineIfVerbose(1, "Boosting iterations: {0}", this.boostingIterations);
			this.PrintLineIfVerbose(1, "Shrinkage: {0}", this.shrinkage);
			this.PrintLineIfVerbose(1, "Loss function: {0}", this.lossFunction);
			this.PrintLineIfVerbose(1, "Proxy loss function: {0}", this.proxyLossFunction);
			this.PrintLineIfVerbose(1, "Proxy loss trimming rate: {0}", this.proxyLossTrimmingRate);
			this.PrintLineIfVerbose(1, "Max tree regions: {0}", this.maxTreeRegions);
			this.PrintLineIfVerbose(1, "Min samples in tree node: {0}", this.minSamplesInNode);
			this.PrintLineIfVerbose(1, "Population size: {0}", this.populationSize);
			this.PrintLineIfVerbose(1, "Max epoch count: {0}", this.maxEpochCount);
			this.PrintLineIfVerbose(1, "Max epoch count without improvement: {0}", this.maxEpochCountWithoutImprovement);
			this.PrintLineIfVerbose(1, "Mutation rate: {0}", this.mutationRate);
			this.PrintLineIfVerbose(1, "Crossover rate: {0}", this.crossoverRate);
			this.PrintLineIfVerbose(1, "Probability of adjust splitter mutation: {0}", this.adjustSplitterMutationProbability);
			this.PrintLineIfVerbose(1, "");
		}

		private Tree<T> CreateBaseTree(ICollection<Sample<T>> trainingSet)
		{
			double mean = this.proxyLossFunction.Minimize(
				new double[trainingSet.Count],
				(from sample in trainingSet select sample.Outcome).ToList());

			return new Tree<T>(mean);
		}

		private Forest<T> PerformBoostingIteration(ICollection<Sample<T>> trainingSet, SplitBase<T> split, Forest<T> previouslyBoostedForest)
		{
			Debug.Assert(trainingSet != null && split != null && previouslyBoostedForest != null);

			List<double> previousOutcomes = CalculateOutcomes(trainingSet, previouslyBoostedForest);
			TreeChromosome<T> ancestor = CreateTreeChromosome(trainingSet, split, previousOutcomes);
			Population<TreeChromosome<T>> population = new Population<TreeChromosome<T>>(ancestor, this.populationSize);
			population.MutationRate = this.mutationRate;
			population.CrossoverRate = this.crossoverRate;

			double lastMax = population.MaxFitness;
			int maxDoNotChangeFor = 0;

			for (int i = 1; i <= maxEpochCount; ++i)
			{
				this.PrintLineIfVerbose(2, "Epoch {0} started", i);

				population.RunEpoch();
				this.CheckTreeRegions(population);

				this.PrintLineIfVerbose(2, "Epoch {0} finished", i);
				this.PrintLineIfVerbose(
					2,
					"Fmax={0:0.0000}, Favg={1:0.0000}",
					population.BestChromosome.Fitness,
					population.AverageFitness);

				if (population.MaxFitness > lastMax)
				{
					lastMax = population.MaxFitness;
					maxDoNotChangeFor = 0;
				}
				else
					maxDoNotChangeFor += 1;

				if (Math.Abs(population.MaxFitness - population.AverageFitness) < 1e-9 ||
					maxDoNotChangeFor > this.maxEpochCountWithoutImprovement)
				{
					this.PrintLineIfVerbose(2, "Genetic algorithm stops early.");
					break;
				}
			}

			this.PrintLineIfVerbose(2, "Mutations happened: {0}", population.MutationsHappened);
			this.PrintLineIfVerbose(2, "Crossovers happened: {0}", population.CrossoversHappened);

			List<Tree<T>> listWithNewTree = new List<Tree<T>> { population.BestChromosome.BuildTree() };
			return new Forest<T>(
				previouslyBoostedForest.BaseTree,
				previouslyBoostedForest.Trees.Concat(listWithNewTree),
				this.shrinkage);
		}

		private static List<double> CalculateOutcomes(ICollection<Sample<T>> samples, Forest<T> forest)
		{
			Debug.Assert(samples != null && samples.Count > 0);
			Debug.Assert(forest != null);

			return (from sample in samples select forest.CalculateOutcome(sample.Properties)).ToList();
		}

		[Conditional("DEBUG")]
		private void CheckTreeRegions(Population<TreeChromosome<T>> population)
		{
			foreach (TreeChromosome<T> chromosome in population.Chromosomes)
				Debug.Assert(chromosome.BuildTree().TreeRegions.Count <= this.maxTreeRegions);
		}

		private TreeChromosome<T> CreateTreeChromosome(ICollection<Sample<T>> trainingSet, SplitBase<T> split, ICollection<double> previousOutcomes)
		{
			return new TreeChromosome<T>(
				split,
				previousOutcomes,
				trainingSet,
				this.proxyLossFunction,
				this.maxTreeRegions,
				this.minSamplesInNode,
				this.adjustSplitterMutationProbability);
		}
	}
}
